package org.finesys.chat.core.base.embedding.store.pgvector.handler;

import java.sql.SQLException;
import java.sql.Statement;
import java.util.Iterator;

import org.finesys.chat.core.base.embedding.store.pgvector.MetadataStorageConfig;

public class JSONBMetadataHandler extends JSONMetadataHandler {
    private final String indexType;

    public JSONBMetadataHandler(MetadataStorageConfig config) {
        super(config);
        if (!this.columnDefinition.getType().equals("jsonb")) {
            throw new IllegalArgumentException("JSONBMetadataHandler requires jsonb type column");
        } else {
            this.indexType = config.indexType();
        }
    }


    @Override
    public void createMetadataIndexes(Statement statement, String index) {
        String indexTypeSql = this.indexType == null ? "" : String.format("USING %s", this.indexType);
        Iterator<String> iterable = this.indexes.iterator();
        while (iterable.hasNext()) {
            String indexName = (String) iterable.next();
            indexName = indexName.trim();
            indexName = this.formatIndex(indexName);
            try {
                String indexSql = String.format("create index if not exists %s_%s on %s %s (%s)", index, indexName, index, indexTypeSql, index);
                statement.executeUpdate(indexSql);
            } catch (SQLException e) {
                throw new RuntimeException(String.format("Error creating index %s on table %s", indexName, index), e);
            }

        }
    }

    private String formatIndex(String index) {
        String indexName = "";
        if (index.contains("->")) {
            String col = this.columnName;
            indexName = String.format("%s_%s", col, index.substring(index.indexOf("->") + 2), index.length() - 1).trim().replaceAll("'", "");
        } else {
            indexName = index.replaceAll(" ", "_");
        }
        return indexName;
    }

}
